﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Security.Cryptography;
using Microsoft.Win32;

namespace System.Peppers.Win32
{
    /// <summary>
    /// Class that allows modification of the driver signing options in windows
    /// *NOTE: this is a registry hack Microsoft would not want us to figure out
    /// </summary>
    public static class DriverSigningSettings
    {
        private static MD5CryptoServiceProvider _md5Provider = new MD5CryptoServiceProvider();

        private static class RegKeys
        {
            public const string
                SEED_KEY = @"HKEY_LOCAL_MACHINE\System\WPA\PnP",
                SEED_VALUE = "Seed",
                DRIVER_SIGNING_KEY = @"HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Driver Signing",
                DRIVER_SIGNING_VALUE = "Policy",
                PRIVATE_HASH_KEY = @"HKEY_LOCAL_MACHINE\Software\Microsoft\Windows\CurrentVersion\Setup",
                PRIVATE_HASH_VALUE = "PrivateHash";
        }

        /// <summary>
        /// Returns the current setting marked valid by windows
        /// NOTE: takes into account the hash value, returning default value if non-match
        /// </summary>
        /// <returns>Current driver signing option</returns>
        public static DriverSigningOption Get()
        {
            DriverSigningOption option = DriverSigningOption.Warn;

            byte[] value = Registry.GetValue(RegKeys.DRIVER_SIGNING_KEY, RegKeys.DRIVER_SIGNING_VALUE, null) as byte[];
            if (value == null || value.Length < 1)
            {
                option = DriverSigningOption.Warn;
            }
            else
            {
                option = (DriverSigningOption)value[0];
            }

            if (option == DriverSigningOption.Warn)
            {
                return DriverSigningOption.Warn;
            }

            byte[] computedHash = CalculateHash(option),
                actualHash = Registry.GetValue(RegKeys.PRIVATE_HASH_KEY, RegKeys.PRIVATE_HASH_VALUE, null) as byte[];

            if (!CompareHashes(computedHash, actualHash))
            {
                return DriverSigningOption.Warn;
            }

            return option;
        }

        /// <summary>
        /// Sets the new value, and calculated hash value so windows counts it valid
        /// </summary>
        /// <param name="value">Value of driver signing option</param>
        public static void Set(DriverSigningOption value)
        {
            Registry.SetValue(RegKeys.PRIVATE_HASH_KEY, RegKeys.PRIVATE_HASH_VALUE, CalculateHash(value));

            Registry.SetValue(RegKeys.DRIVER_SIGNING_KEY, RegKeys.DRIVER_SIGNING_VALUE, new byte[] { (byte)value });
        }

        private static byte[] CalculateHash(DriverSigningOption value)
        {
            int seed = (int)Registry.GetValue(RegKeys.SEED_KEY, RegKeys.SEED_VALUE, 0);

            byte[] seedBytes = BitConverter.GetBytes(seed);

            byte[] input = new byte[] 
                { 
                    0, (byte)value, 0, 0, 
                    seedBytes[0],
                    seedBytes[1],
                    seedBytes[2],
                    seedBytes[3],
                };

            return _md5Provider.ComputeHash(input);
        }

        /// <summary>
        /// Compares 2 byte arrays, might be a method already in .Net?  I know Linq has one
        /// </summary>
        /// <param name="hash1">byte array to compare</param>
        /// <param name="hash2">byte array to compare</param>
        /// <returns>True if equal</returns>
        private static bool CompareHashes(byte[] hash1, byte[] hash2)
        {
            if (hash1 == null || hash2 == null || hash1.Length != hash2.Length)
            {
                return false;
            }
            
            int length = hash1.Length;
            for (int i = 0; i < length; i++)
            {
                if (hash1[i] != hash2[i])
                {
                    return false;
                }
            }

            return true;
        }
    }
}
